import os
import numpy as np
import pandas as pd
import pickle
from gensim.models import word2vec
import threading


def train():
    print("### start word2vec processing ###")
    train_list_pos = "/home/datanfs/macong_data/tencent_data/train_preliminary/train_data/h5word2vec.pkl"
    train_store_pos = "/home/datanfs/macong_data/tencent_data/train_preliminary/train_data/word2vec_64.wv"
    with open(train_list_pos, "rb") as f:
        train_data = pickle.load(f)
        train_data = train_data["data"]
        trans_int2str = []
        print("### start translate int 2 str ###")
        for x in train_data:
            trans_int2str.append(list(map(str, x)))
        print("### start translate int 2 str over ###")
        workers = os.cpu_count()
        print("### {} threads totally ###".format(workers))
        model = word2vec.Word2Vec(trans_int2str, size=64, min_count=2, window=10, iter=30, workers=workers)
        model.save(train_store_pos)

if __name__ == '__main__':
    train()

